import requests
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
IPADDRESS = "localhost" #Replace this with the correct IP address
UNSTRUCTURED_DATA_PORT = "8086"
# Helper function to ingest document in canonical RAG retriever
def ingest_file(filepath: str) -> bool:
    """
    Ingest file in canonical RAG retriever

    Args:
        filepath: Path to the file to be ingested in retreiver

    Returns:
        bool: Status of file ingestion
    """
    # URL of the API endpoint
    url = f'http://{IPADDRESS}:{UNSTRUCTURED_DATA_PORT}/documents'

    # Open the file in binary mode and send it in a POST request
    with open(filepath, 'rb') as file:
        files = {'file': file}
        try:
            response = requests.post(url, files=files)
            return response.status_code == 200
        except requests.exceptions.RequestException as e:
            print(f"Request failed for {filepath}: {e}")
            return False
#%%


directory_path = '../data/product'
max_workers = 5  # Adjust this based on your system's capabilities and API limits

filepaths = [os.path.join(directory_path, filename) for filename in os.listdir(directory_path) if filename.endswith(".txt")]

successfully_ingested = []
failed_ingestion = []

with ThreadPoolExecutor(max_workers=max_workers) as executor:
    future_to_file = {executor.submit(ingest_file, filepath): filepath for filepath in filepaths}

    for future in as_completed(future_to_file):
        filepath = future_to_file[future]
        try:
            if future.result():
                print(f"Successfully Ingested {os.path.basename(filepath)}")
                successfully_ingested.append(filepath)
            else:
                print(f"Failed to Ingest {os.path.basename(filepath)}")
                failed_ingestion.append(filepath)
        except Exception as e:
            print(f"Exception occurred while ingesting {os.path.basename(filepath)}: {e}")
            # traceback.print_exc()
            failed_ingestion.append(filepath)

print(f"Total files successfully ingested: {len(successfully_ingested)}")
print(f"Total files failed ingestion: {len(failed_ingestion)}")